import torch.nn as nn

from models.backbone.resnet import resnet18


class Segmentation_Model(nn.Module):
    def __init__(self):
        super(Segmentation_Model, self).__init__()
        self.conv4_0 = nn.Sequential(
            nn.Conv2d(512, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv4_1 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv3_0 = nn.Sequential(
            nn.Conv2d(256, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv3_1 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv2_0 = nn.Sequential(
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv2_1 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv1_0 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.conv1_1 = nn.Sequential(
            nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(1)
        )
        self.init_weights()
        self.backbone = resnet18(pretrained=True)
        self.up = nn.UpsamplingBilinear2d(scale_factor=2)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, img):
        x4, x3, x2, x1 = self.backbone(img)
        x4 = self.conv4_0(x4)
        x3 = self.up(self.conv4_1(x4)) + self.conv3_0(x3)
        x2 = self.up(self.conv3_1(x3)) + self.conv2_0(x2)
        x1 = self.up(self.conv2_1(x2)) + self.conv1_0(x1)
        x = self.conv1_1(x1).sigmoid()
        return x
